import numpy as np
import pandas as pd
import os
from sklearn.tree import DecisionTreeRegressor
from sklearn.preprocessing import LabelEncoder
from sklearn.externals import joblib
import operator
from functools import reduce


class DecisionTree_train(object):
    def __init__(self, data):
        self.data = data
        self.script_path = (os.path.dirname(os.path.realpath(__file__)))
        self.filepath = os.path.join(
            self.script_path, 'model', 'Decision_Tree.m')
        self.le = LabelEncoder()
        self.fit_data(data)

    #保存模型
    def Save_Model(self, model, filepath):
        joblib.dump(model, filename=filepath)

    #加载模型
    def Load_Model(self, filepath):
        model = joblib.load(filepath)
        return model

    def Decision_Tree_Regressor(self, data):
        x_train, y_train = data

        try:
            reg = self.Load_Model(self.filepath)
        except Exception:
            reg = DecisionTreeRegressor(criterion='mse')
            print("创建新的训练模型!")

        #训练模型
        reg.fit(x_train, y_train)

        self.Save_Model(reg, self.filepath)
        return reg

    def fit_data(self, data):

        if isinstance(data, pd.DataFrame):
            pass
        elif isinstance(data, list):
            data = pd.DataFrame(data, columns=[
                                'id', 'cve', 'version', 'vul', 'vendor', 'product_type', 'product'])

        data = np.array(data).tolist()
        self.le.fit(reduce(operator.add, data))
    def predict_result(self, x_test, filepath=None):
        if filepath == None:
            filepath = self.filepath

        data = pd.DataFrame(x_test, columns=['product'])
        x_test = self.transform_form(data).reshape(-1, 1)
        reg = self.Load_Model(filepath)

        #预测结果
        predict_y = reg.predict(x_test)
        predict_int_y = list(map(int, predict_y[:]))
        predict_y = self.le.inverse_transform(predict_int_y)

        return predict_y

    def data_deal(self, data):

        if isinstance(data, pd.DataFrame):
            pass
        elif isinstance(data, list):
            data = pd.DataFrame(data, columns=[
                                'id', 'cve', 'version', 'vul', 'vendor', 'product_type', 'product'])
        x = data.iloc[:, 6]
        y = data.iloc[:, 1]
        x = self.transform_form(x).reshape(-1, 1)
        y = self.transform_form(y).reshape(-1, 1)

        data = [x, y]
        return data

    def transform_form(self, data):
        xx = data.values.tolist()
        # self.le.fit(xx)
        yy = self.le.transform(xx)
        return yy

    def train_result(self, train_data):
        data = self.data_deal(train_data)
        result = self.Decision_Tree_Regressor(data)
        return result

    def get_inf(self, output):
        os = []
        for i in range(len(output)):
            try:
                print(output[i]["os"][0]["osclass"]["vendor"])
            except:
                os.append("Mac Os X")
        print("===============")
        print(os)
        print("===============")
        return os


if __name__ == "__main__":

    data = pd.read_csv("os_vul.csv")
    data1 = [
        ["10301", "CVE-2018-12015", "-",
            """Directory traversal"",""Bypass a restriction or similar""", "Apple", "OS", "Mac Os X"],
        ["10302", "CVE-2018-12015", "10.0", """Directory traversal"",""Bypass a restriction or similar""", "Apple", "OS", "Mac Os X"]]

    data2 = ["Mac Os X"]
    dt = DecisionTree_train(data)
    print(dt.train_result(data))
    print(dt.predict_result(data2))
